import os
from openai import OpenAI
from dotenv import load_dotenv
load_dotenv()   # 加载环境变量(有缓存)

class LLMClient:
    def __init__(self):
        
        # 环境变量的获取
        # self.xxx 表示给该类添加字段
        self.openai_api_key = os.getenv('DASHSCOPE_API_KEY') or os.getenv("API_KEY")
        self.model = os.getenv("MODEL") or ""
        self.base_url = os.getenv("BASE_URL") 

        # 检查环境变量是否为空
        if not self.openai_api_key:
            raise ValueError("DASHSCOPE_API_KEY 环境变量未获取成功")
        if not self.base_url:
            raise ValueError("BASE_URL 环境变量未获取成功")

        # 初始化OpenAI
        self.client = OpenAI(api_key=self.openai_api_key,base_url=self.base_url)

    
    def chat_loop(self):
        """
        进入聊天循环，等待用户输入并处理。

        该方法初始化提示信息，进入一个无限循环等待用户输入。用户可以输入问题，系统会调用
        process_query 方法处理问题并返回结果。如果用户输入 "quit"，则退出循环。

        具体步骤如下：
        1. 打印初始化提示信息，告知用户客户端已启动，并提示如何退出。
        2. 进入一个无限循环，等待用户输入。
        3. 获取用户输入，并去除首尾空白字符。
        4. 检查用户输入是否为 "quit"，如果是则退出循环。
        5. 调用 process_query 方法处理用户输入，并获取响应。
        6. 打印 AI 的响应。
        7. 如果在处理过程中发生异常，捕获异常并打印错误信息。

        :return: 无返回值
        """

        # 初始化提示信息
        print("\n🤖 MCP 客户端已启动！输入 'quit' 退出")
        # 进入主循环中等待用户输入
        while True:
            try:
                query = input("\n你: ").strip()
                if query.lower() == "quit":
                    break

                self.process_query_stream(query)
            except Exception as e:
                print(f"\n⚠ 发生错误: {str(e)}")

    # 非流式调用
    def process_query(self, query:str):
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "你是一个很厉害的助手"},
                {"role": "user", "content": f"{query}"},
            ]
        )
        print(f"\n🤖 AI: {response.choices[0].message.content}")
    
    # 流式调用
    def process_query_stream(self, query:str):
         response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "你是一个很厉害的助手"},
                {"role": "user", "content": f"{query}"},
            ],
            stream=True,
            stream_options={"include_usage": True} # 显示token数量
        )
         print("\n🤖 AI:")
         for chunk in response:
             if chunk.choices:
                print(chunk.choices[0].delta.content,end="")
             else:
                if chunk.usage:
                    print("")
                    print(f"total_tokens:{chunk.usage.total_tokens}")
    
    

def main():
    mcpClient = LLMClient()
    # 进入 chat loop
    mcpClient.chat_loop() 

if __name__ == '__main__':
    main()